# document_upload_api.py
import os
import tempfile
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from dotenv import load_dotenv
from pymilvus import connections, Collection, utility
import docx2txt
import PyPDF2
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
import dashscope
import chardet

# 加载环境变量
load_dotenv()

# 初始化FastAPI应用
app = FastAPI(title="文档上传API", description="上传文档并存储到Milvus向量数据库")

# 设置通义千问API密钥
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')

# 配置Milvus连接
milvus_host = os.getenv('MILVUS_HOST', 'localhost')
milvus_port = os.getenv('MILVUS_PORT', '19530')

# 尝试连接Milvus服务器
try:
    connections.connect(
        host=milvus_host,
        port=milvus_port
    )
    print(f"成功连接到Milvus服务器: {milvus_host}:{milvus_port}")
except Exception as e:
    print(f"无法连接到Milvus服务器: {str(e)}")
    # 不抛出异常，而是在API请求时处理

# 定义集合名称和向量维度
collection_name = os.getenv('MILVUS_COLLECTION_NAME', 'contract_documents_v1')
dimension = 1536  # 通义千问text-embedding-v2模型的维度

# 初始化集合变量
collection = None

try:
    # 检查集合是否存在
    if utility.has_collection(collection_name):
        collection = Collection(collection_name)
        collection.load()
        print(f"成功加载Milvus集合: {collection_name}")
    else:
        print(f"Milvus集合 '{collection_name}' 不存在")
except Exception as e:
    print(f"加载Milvus集合时出错: {str(e)}")

# 初始化文本分割器
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=800,
    chunk_overlap=100
)

def get_embedding(text: str, model="text-embedding-v2"):
    """调用通义千问的Embedding API生成向量"""
    try:
        if not text.strip():
            print("文本内容为空，无法生成Embedding")
            return None
        
        from dashscope import TextEmbedding
        response = TextEmbedding.call(
            model=model,
            input=[text]
        )
        
        if not response or not response.output or not response.output['embeddings']:
            print("Embedding API返回空结果")
            return None
        
        return response.output['embeddings'][0]['embedding']
    except Exception as e:
        print(f"生成Embedding时出错: {str(e)}")
        return None

async def process_uploaded_file(file: UploadFile) -> list:
    """处理上传的文件并分割成文本块"""
    try:
        file_name = file.filename
        ext = os.path.splitext(file_name)[1].lower()
        text = ""
        
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file:
            temp_file.write(await file.read())
            temp_file_path = temp_file.name
        
        # 根据文件类型处理
        if ext == '.pdf':
            with open(temp_file_path, 'rb') as f:
                reader = PyPDF2.PdfReader(f)
                for page in reader.pages:
                    page_text = page.extract_text() or ""
                    text += page_text
        elif ext == '.docx':
            text = docx2txt.process(temp_file_path)
        elif ext == '.doc':
            # 对于.doc文件，尝试使用pywin32处理 (Windows专用)
            try:
                import win32com.client
                print("使用pywin32处理.doc文件...")
                word = win32com.client.Dispatch("Word.Application")
                doc = word.Documents.Open(temp_file_path)
                text = doc.Content.Text
                doc.Close()
                word.Quit()
                print("成功使用pywin32提取.doc文件内容")
            except ImportError:
                raise HTTPException(status_code=501, detail="处理.doc文件需要安装pywin32库。您可以通过 'pip install pywin32' 安装，或者将.doc文件转换为.docx格式后再上传")
            except Exception as e:
                print(f"处理.doc文件时出错: {str(e)}")
                raise HTTPException(status_code=400, detail=f"无法解析.doc文件: {str(e)}")
        elif ext in ['.txt', '.md']:
            # 检测文件编码
            with open(temp_file_path, 'rb') as f:
                raw_data = f.read(10000)  # 读取前10000字节进行检测
                result = chardet.detect(raw_data)
                detected_encoding = result['encoding']
                confidence = result['confidence']
                
                # 如果检测到的编码置信度较低，尝试常见编码
                if not detected_encoding or confidence < 0.7:
                    print(f"编码检测不确定(置信度: {confidence})，尝试常见编码...")
                    common_encodings = ['utf-8', 'utf-16le', 'gbk', 'latin1']
                    for encoding in common_encodings:
                        try:
                            with open(temp_file_path, 'r', encoding=encoding) as f_enc:
                                text = f_enc.read()
                            detected_encoding = encoding
                            print(f"成功使用 {encoding} 编码读取文件")
                            break
                        except UnicodeDecodeError:
                            continue
                
                # 使用检测到的编码读取文件
                if detected_encoding:
                    try:
                        with open(temp_file_path, 'r', encoding=detected_encoding) as f:
                            text = f.read()
                        print(f"使用检测到的编码 {detected_encoding} 读取文件成功")
                    except Exception as e:
                        raise HTTPException(status_code=400, detail=f"无法使用编码 {detected_encoding} 读取文件: {str(e)}")
                else:
                    raise HTTPException(status_code=400, detail="无法确定文件编码")
        else:
            # 尝试作为文本文件读取
            try:
                # 检测编码
                with open(temp_file_path, 'rb') as f:
                    raw_data = f.read(10000)
                    result = chardet.detect(raw_data)
                    detected_encoding = result['encoding']
                    
                if detected_encoding:
                    with open(temp_file_path, 'r', encoding=detected_encoding) as f:
                        text = f.read()
                else:
                    raise HTTPException(status_code=400, detail="无法确定文件编码")
            except Exception as e:
                raise HTTPException(status_code=400, detail=f"不支持的文件类型或无法解析文件: {str(e)}")
        
        # 删除临时文件
        os.unlink(temp_file_path)
        
        # 创建Document对象并分割
        if text.strip():
            document = Document(page_content=text, metadata={'file_name': file_name})
            chunks = text_splitter.split_documents([document])
            return chunks
        else:
            raise HTTPException(status_code=400, detail=f"上传的文件内容为空")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理文件时出错: {str(e)}")

def store_in_milvus(chunks: list) -> dict:
    """将处理好的文本块向量化并存入Milvus"""
    file_names = []
    text_chunks = []
    vectors = []
    success_count = 0
    fail_count = 0
    
    for chunk in chunks:
        text = chunk.page_content
        vector = get_embedding(text)
        
        if vector is not None:
            file_names.append(chunk.metadata['file_name'])
            text_chunks.append(text)
            vectors.append(vector)
            success_count += 1
        else:
            fail_count += 1
    
    # 批量插入数据
    if vectors:
        try:
            data = [file_names, text_chunks, vectors]
            insert_result = collection.insert(data)
            collection.flush()  # 刷新使数据可搜索
            return {
                "success_count": success_count,
                "fail_count": fail_count,
                "inserted_ids": list(insert_result.primary_keys)
            }
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"向Milvus插入数据时出错: {str(e)}")
    else:
        return {
            "success_count": 0,
            "fail_count": fail_count,
            "inserted_ids": []
        }

@app.post("/upload-document/")
@app.post("/upload-document")
async def upload_document(file: UploadFile = File(...)):
    """
    上传文档并存储到Milvus向量数据库
    - 支持的文件类型: PDF, DOCX, TXT, MD
    - 文件会被分割成文本块并向量化
    - 向量和文本块会被存储到Milvus
    """
    # 检查Milvus连接和集合状态
    if collection is None:
        raise HTTPException(status_code=503, detail="无法连接到Milvus数据库或集合未加载，请检查Milvus服务状态")
    
    # 检查文件类型
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in ['.pdf', '.docx', '.doc', '.txt', '.md']:
        raise HTTPException(status_code=400, detail=f"不支持的文件类型: {ext}，请上传PDF, DOCX, TXT或MD文件")
    
    # 处理上传的文件
    chunks = await process_uploaded_file(file)
    
    # 存储到Milvus
    result = store_in_milvus(chunks)
    
    return JSONResponse(content={
        "message": "文件上传处理成功",
        "file_name": file.filename,
        "chunks_count": len(chunks),
        "result": result
    })

@app.get("/")
def root():
    return {
        "message": "文档上传API服务正常运行，请使用POST /upload-document/接口上传文档"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8002)